use std::sync::Arc;

use axum::{
    Extension,
    extract::{
        WebSocketUpgrade,
        ws::{Message, Utf8Bytes},
    },
    response::IntoResponse,
};
use dashmap::{DashMap, DashSet};
use futures::{Sink, SinkExt, Stream, StreamExt};
use tokio::sync::broadcast;

use super::{Msg, MsgData};

#[derive(Debug)]
struct StateInner {
    // 同一room下的users
    room_users: DashMap<String, DashSet<String>>,
    // 一个user订阅的rooms
    user_rooms: DashMap<String, DashSet<String>>,

    tx: broadcast::Sender<Arc<Msg>>,
}

// 需作为Extension共享, 通过Arc实现clone
#[derive(Debug, Clone)]
pub struct State(Arc<StateInner>);

const CAPACITY: usize = 64;

impl Default for StateInner {
    fn default() -> Self {
        let (tx, _rx) = broadcast::channel(CAPACITY);
        Self {
            room_users: Default::default(), // dashmap实现了default
            user_rooms: Default::default(),
            tx,
        }
    }
}

impl State {
    // 可用default trait替代
    pub fn new() -> Self {
        Self(Arc::new(StateInner::default()))
    }

    pub fn get_rooms_by_user(&self, user_name: &str) -> Vec<String> {
        self.0
            .user_rooms
            .get(user_name)
            .map(|rooms| rooms.clone().into_iter().collect())
            .unwrap_or_default()
    }
}

pub async fn ws_handle(
    ws: WebSocketUpgrade,
    Extension(state): Extension<State>,
) -> impl IntoResponse {
    //ws.on_upgrade(handler);
    ws.on_upgrade(|socket| handle(socket, state))
}

async fn handle<S>(socket: S, state: State)
where
    // 这里 'static 表示socket对象拥有所有权
    S: Stream<Item = Result<Message, axum::Error>> + Sink<Message> + Send + 'static,
{
    // 聊天消息的订阅端
    let mut rx = state.0.tx.subscribe();
    // socket接收发送端
    let (mut sender, mut receiver) = socket.split();

    let state_inner = state.clone().0;

    // 当一个socket连上后, 创建两个协程进行收发消息
    let mut rece_task = tokio::spawn(async move {
        while let Some(Ok(data)) = receiver.next().await {
            match data {
                // websocket中提供事件类型
                Message::Text(msg) => {
                    handle_message(msg.as_str().try_into().unwrap(), state_inner.clone()).await
                }
                _ => {}
            }
        }
    });

    let mut send_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            let data: String = msg.as_ref().try_into().unwrap();
            if sender
                .send(Message::Text(Utf8Bytes::from(data)))
                .await
                .is_err()
            {
                tracing::warn!("fail to send msg to client");
                return;
            };
        }
    });

    // 如果有一个任务退出, 则需要关闭另一个任务
    tokio::select! {
        _ = &mut rece_task => send_task.abort(),
        _ = &mut send_task => rece_task.abort(),
    }
    tracing::info!("websocket disconnected");

    // 清理state中的状态
    let user_name = "todo_user_name";
    for room in state.get_rooms_by_user(user_name) {
        if let Err(e) = state
            .0
            .tx
            .send(Arc::new(Msg::leave(room.as_str(), user_name)))
        {
            tracing::warn!("fail to send leave msg: {e}");
        }
    }
}

// TODO: 直接传入Arc
async fn handle_message(msg: Msg, state: Arc<StateInner>) {
    let msg = match msg.data {
        MsgData::Join => {
            // 找到聊天室加入
            state
                .room_users
                .entry(msg.room.clone())
                .or_default()
                .insert(msg.user_name.clone());
            // 记录用户加入了哪些聊天室
            state
                .user_rooms
                .entry(msg.user_name.clone())
                .or_insert_with(DashSet::new)
                .insert(msg.room.clone());
            msg
        }
        MsgData::Leave => {
            if let Some(room) = state.room_users.get_mut(&msg.room) {
                room.remove(&msg.user_name);
                if room.is_empty() {
                    state.room_users.remove(&msg.room);
                }
            }
            if let Some(rooms) = state.user_rooms.get_mut(&msg.user_name) {
                rooms.remove(&msg.room);
                if rooms.is_empty() {
                    state.user_rooms.remove(&msg.user_name);
                }
            }
            msg
        }
        _ => msg,
    };

    // 将msg发送给所有订阅者
    if let Err(e) = state.tx.send(Arc::new(msg)) {
        tracing::warn!("fail to send msg: {e}");
    }
}

// echo
/* while let Some(msg) = socket.recv().await {
    let msg = if let Ok(msg) = msg {
        msg
    } else {
        // client disconnected
        return;
    };

    if socket.send(msg).await.is_err() {
        // client disconnected
        return;
    }
} */
